import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict, Union

from rainforeLearn.gomoku.v2.neywork.blocks import ResidualBlock, NoisyLinear
from rainforeLearn.gomoku.v2.neywork.output_heads import NetworkOutputHeads
from rainforeLearn.gomoku.v2.neywork.positional_encoding import LearnablePositionalEncoding2D, PositionalEncoding2D


class GomokuNetwork(nn.Module):
    """增强版五子棋神经网络 - 支持多任务输出"""
    
    def __init__(self, board_size: int = 15, num_layers: int = 16, num_channels: int = 256,
                 use_noisy: bool = False, dropout_rate: float = 0.1, history_steps: int = 4,
                 attention_type: str = 'cbam', attention_freq: int = 4,
                 task_type: str = 'classification', num_classes: Optional[int] = None,
                 use_positional_encoding: bool = True,
                 learnable_pos_encoding: bool = True,
                 pos_encoding_temperature: float = 10000.0):
        super().__init__()
        
        # 基本配置
        self.board_size = board_size
        self.num_channels = num_channels
        self.use_noisy = use_noisy
        self.history_steps = history_steps
        self.task_type = task_type
        self.use_positional_encoding = use_positional_encoding
        
        # 确定输出维度
        self._setup_output_dimensions(task_type, num_classes, board_size)
        
        # 构建主干网络
        input_channels = 2 + history_steps
        
        # 添加位置编码模块
        if self.use_positional_encoding:
            
            if learnable_pos_encoding:
                self.pos_encoding = LearnablePositionalEncoding2D(
                    input_channels, board_size, board_size
                )
            else:
                self.pos_encoding = PositionalEncoding2D(
                    input_channels, board_size, board_size, 
                    temperature=pos_encoding_temperature
                )
        else:
            self.pos_encoding = None
            
        self.backbone = self._build_backbone(input_channels, num_channels, num_layers, 
                                           attention_type, attention_freq)
        
        # 构建输出头
        self._build_output_heads(dropout_rate)
        
        self.relu = nn.ReLU(inplace=True)
        print(f"网络参数量: {self.count_parameters():,}")
    
    def _setup_output_dimensions(self, task_type: str, num_classes: Optional[int], board_size: int) -> None:
        """设置输出维度"""
        if task_type == 'classification':
            self.output_size = num_classes or (board_size * board_size)
        elif task_type == 'regression':
            self.output_size = num_classes or 1
        elif task_type == 'multi_task':
            self.action_size = board_size * board_size
            self.value_size = 1
        else:
            raise ValueError(f"Unsupported task_type: {task_type}")
    
    def _build_backbone(self, input_channels: int, num_channels: int, num_layers: int,
                       attention_type: str, attention_freq: int) -> nn.Sequential:
        """构建主干网络"""
        layers = []
        
        # 第一层：输入通道到隐藏通道
        first_attention = attention_type if 0 % attention_freq == 0 else None
        layers.append(ResidualBlock(input_channels, num_channels, use_attention=first_attention))
        
        # 中间层：隐藏通道到隐藏通道
        for i in range(1, num_layers):
            use_attention = attention_type if i % attention_freq == 0 else None
            layers.append(ResidualBlock(num_channels, num_channels, use_attention=use_attention))
        
        return nn.Sequential(*layers)
    
    def _build_output_heads(self, dropout_rate: float) -> None:
        """构建输出头"""
        if self.task_type == 'multi_task':
            NetworkOutputHeads.build_multi_task_heads(self, dropout_rate)
        else:
            NetworkOutputHeads.build_single_task_head(self, dropout_rate)
    
    # ========================================================================
    # 前向传播方法
    # ========================================================================
    
    def forward(self, state: torch.Tensor) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
        """前向传播"""
        # 处理输入维度
        state = self._preprocess_input(state)
        
        # 应用位置编码
        if self.pos_encoding is not None:
            state = self.pos_encoding(state)

        # 特征提取
        features = self.backbone(state)
        
        # 根据任务类型进行输出
        if self.task_type == 'multi_task':
            return self._forward_multi_task(features)
        else:
            return self._forward_single_task(features)
    
    def _preprocess_input(self, state: torch.Tensor) -> torch.Tensor:
        """预处理输入状态"""
        if len(state.shape) == 2:
            batch_size = state.shape[0]
            state = state.view(batch_size, self.board_size, self.board_size).unsqueeze(1)
            zeros = torch.zeros_like(state)
            additional_channels = [zeros] + [zeros.clone() for _ in range(self.history_steps)]
            state = torch.cat([state] + additional_channels, dim=1)
        return state
    
    def _forward_multi_task(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
        """多任务前向传播"""
        action_logits = self._forward_action_branch(features)
        value_pred = self._forward_value_branch(features)
        
        return {
            'action_logits': action_logits,
            'value_pred': value_pred
        }
    
    def _forward_action_branch(self, features: torch.Tensor) -> torch.Tensor:
        """动作预测分支前向传播"""
        out = self.action_conv(features)
        out = self.relu(self.action_bn(out))
        out = self.action_dropout(out)
        
        out = out.view(out.size(0), -1)
        out = self.relu(self.action_fc1(out))
        out = self.relu(self.action_fc2(out))
        
        return self.action_fc3(out)
    
    def _forward_value_branch(self, features: torch.Tensor) -> torch.Tensor:
        """价值评估分支前向传播"""
        out = self.value_conv(features)
        out = self.relu(self.value_bn(out))
        out = self.value_dropout(out)
        
        out = out.view(out.size(0), -1)
        out = self.relu(self.value_fc1(out))
        out = self.relu(self.value_fc2(out))
        
        return self.value_fc3(out)
    
    def _forward_single_task(self, features: torch.Tensor) -> torch.Tensor:
        """单任务前向传播"""
        out = self.output_conv(features)
        out = self.relu(self.output_bn(out))
        out = self.output_dropout(out)
        
        out = out.view(out.size(0), -1)
        out = self.relu(self.output_fc1(out))
        out = self.relu(self.output_fc2(out))
        output = self.output_fc3(out)
        
        if self.task_type == 'classification':
            return F.log_softmax(output, dim=1)
        else:  # regression
            return output
    
    # ========================================================================
    # 工具方法
    # ========================================================================
    
    def reset_noise(self) -> None:
        """重置所有NoisyNet层的噪声"""
        if self.use_noisy:
            for module in self.modules():
                if isinstance(module, NoisyLinear):
                    module.reset_noise()
    
    def count_parameters(self) -> int:
        """计算模型参数量"""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)